package club.drguo.hadoop.mapreduce.enhance;

import java.io.IOException;

import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.mapreduce.RecordWriter;
import org.apache.hadoop.mapreduce.TaskAttemptContext;
import org.apache.hadoop.mapreduce.lib.output.FileOutputFormat;

public class EnhanceOutputFormat<K, V> extends FileOutputFormat<K, V>{

	@Override
	public RecordWriter<K, V> getRecordWriter(TaskAttemptContext arg0) throws IOException, InterruptedException {
		FileSystem fs = FileSystem.get(new Configuration());
		FSDataOutputStream enhanceOS = fs.create(new Path("/flow/output/enhanceLog"));
		FSDataOutputStream tocrawlOS = fs.create(new Path("/flow/output/tocrawl"));
		return new EnhanceRecordWriter<K, V>(enhanceOS, tocrawlOS);
	}
	public static class EnhanceRecordWriter<K, V> extends RecordWriter<K, V>{
		private FSDataOutputStream enhanceOS = null;
		private FSDataOutputStream tocrawlOS  = null; 
		public EnhanceRecordWriter(FSDataOutputStream enhanceOS, FSDataOutputStream tocrawlOS) {
			this.enhanceOS = enhanceOS;
			this.tocrawlOS = tocrawlOS;
		}
		@Override
		public void write(K key, V value) throws IOException, InterruptedException {
			
			if(key.toString().endsWith("tocrawl")){
				tocrawlOS.write(key.toString().getBytes());
			}else{
				enhanceOS.write(key.toString().getBytes());
			}
		}
		@Override
		public void close(TaskAttemptContext context) throws IOException, InterruptedException {
			if(enhanceOS != null){
				enhanceOS.close();
			}
			if(tocrawlOS != null){
				tocrawlOS.close();
			}
		}

		
	}
}
